// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/testing/exec/presto_query_runner_timestamp_with_time_zone_transform.h>
#include <pollux/functions/lib/date_time_formatter.h>
#include <pollux/functions/prestosql/types/timestamp_with_time_zone_type.h>
#include <pollux/parse/expressions.h>
#include <pollux/type/tz/time_zone_map.h>
#include <pollux/vector/simple_vector.h>

namespace kumo::pollux::exec::test {
namespace {
const std::string kFormat = "yyyy-MM-dd HH:mm:ss.SSS ZZZ";
const std::string kBackupFormat = "yyyy-MM-dd HH:mm:ss.SSS ZZ";

std::string format(const int64_t timestampWithTimeZone) {
  static const std::shared_ptr<functions::DateTimeFormatter> kJodaDateTime =
      functions::buildJodaDateTimeFormatter(kFormat).value();

  const auto timestamp = unpackTimestampUtc(timestampWithTimeZone);
  const auto timeZoneId = unpackZoneKeyId(timestampWithTimeZone);
  auto* timezonePtr = tz::locateZone(tz::getTimeZoneName(timeZoneId));

  const auto maxResultSize = kJodaDateTime->maxResultSize(timezonePtr);
  std::string str;
  str.resize(maxResultSize);
  const auto resultSize =
      kJodaDateTime->format(timestamp, timezonePtr, maxResultSize, str.data());
  str.resize(resultSize);

  return str;
}
} // namespace

// Convert a TimestampWithTimeZone into a Varchar using DatetimeFormatter, so
// that we can get the TimestampWithTimeZone back by calling parse_datetime.
variant TimestampWithTimeZoneTransform::transform(
    const BaseVector* const vector,
    vector_size_t row) const {
  POLLUX_CHECK(isTimestampWithTimeZoneType(vector->type()));
  POLLUX_CHECK(!vector->is_null_at(row));

  return variant::create<TypeKind::VARCHAR>(
      format(vector->as_checked<SimpleVector<int64_t>>()->value_at(row)));
}

// Applies parse_datetime to a Vector of VARCHAR (formatted timestamps with time
// zone) to produce values of type TimestampWithTimeZone.
core::ExprPtr TimestampWithTimeZoneTransform::projectionExpr(
    const core::ExprPtr& inputExpr,
    const std::string& columnAlias) const {
  // format_datetime with the ZZZ pattern produces time zones that need to be
  // parsed with either the ZZZ or ZZ pattern, to handle this we try parsing
  // with the ZZZ pattern first, and then the ZZ pattern if that fails.
  // coalesce(try(parse_datetime(..., '... ZZZ')), parse_datetime(..., '...
  // ZZ'))
  return std::make_shared<core::CallExpr>(
      "coalesce",
      std::vector<core::ExprPtr>{
          std::make_shared<core::CallExpr>(
              "try",
              std::vector<core::ExprPtr>{std::make_shared<core::CallExpr>(
                  "parse_datetime",
                  std::vector<core::ExprPtr>{
                      inputExpr,
                      std::make_shared<core::ConstantExpr>(
                          VARCHAR(),
                          variant::create<TypeKind::VARCHAR>(kFormat),
                          std::nullopt)},
                  std::nullopt)},
              std::nullopt),
          std::make_shared<core::CallExpr>(
              "parse_datetime",
              std::vector<core::ExprPtr>{
                  inputExpr,
                  std::make_shared<core::ConstantExpr>(
                      VARCHAR(),
                      variant::create<TypeKind::VARCHAR>(kBackupFormat),
                      std::nullopt)},
              std::nullopt)},
      columnAlias);
}
} // namespace kumo::pollux::exec::test
